import pandas as pd
from sklearn.metrics import mean_absolute_error
from catboost import CatBoostRegressor
from sklearn.model_selection import GridSearchCV

# 读取数据
data = pd.read_csv('data.csv')

# 将 time 列转换为 datetime 类型
data['time'] = pd.to_datetime(data['time'])

# 确保数据按时间排序
data = data.sort_values(by='time').reset_index(drop=True)

# 创建一个新的列，存储三分钟前的IC_DCS_CHIMNEY_NOX值
# 注意：这会导致数据的前三行没有对应的“三分钟前”的值，因此我们需要删除这些行
data['IC_DCS_CHIMNEY_NOX_lagged'] = data['IC_DCS_CHIMNEY_NOX'].shift(3)
data = data.dropna()  # 删除包含NaN的行

# 现在IC_DCS_CHIMNEY_NOX_lagged是目标变量
y = data['IC_DCS_CHIMNEY_NOX_lagged']

# 从 time 列中提取分钟（或其他时间特征，如果需要）
# 但注意，这里我们不再直接使用minute作为特征，因为我们已经确保了时间顺序
# 并且将使用整个时间序列来划分训练集和测试集
# X = data[['minute', ...]]  # 如果需要其他时间特征，可以在这里添加
X = data.drop(columns=['IC_DCS_CHIMNEY_NOX', 'IC_DCS_CHIMNEY_NOX_lagged', 'time'])  # 剩余列作为特征

# 基于时间戳划分训练集和测试集
# 我们将使用一个固定的时间点来划分，而不是随机划分
# 假设我们使用最后20%的时间段作为测试集
train_cutoff = data['time'].quantile(0.8)
X_train = X[data['time'] < train_cutoff]
y_train = y[data['time'] < train_cutoff]
X_test = X[data['time'] >= train_cutoff]
y_test = y[data['time'] >= train_cutoff]

# 为了绘制时间序列图，我们需要保留时间戳
time_train = data['time'][data['time'] < train_cutoff]
time_test = data['time'][data['time'] >= train_cutoff]

# 初始化CatBoostRegressor模型
model = CatBoostRegressor(iterations=1000,
                             learning_rate=0.1,
                             depth=6,
                             loss_function='MAE',
                             verbose=False)

# 定义参数网格
param_grid = {
    'iterations': [1500, 2000, 2500],
    'learning_rate': [0.005, 0.01, 0.1],
    'depth': [1, 2, 4],
    'l2_leaf_reg': [1, 3, 5]  # 可以添加其他您认为重要的参数
}

# 初始化CatBoostRegressor模型（注意：这里我们不设置具体的参数值，因为将在表格搜索中确定）
model_base = CatBoostRegressor(loss_function='MAE', verbose=False)

# 初始化GridSearchCV
grid_search = GridSearchCV(estimator=model_base, param_grid=param_grid, scoring='neg_mean_absolute_error', cv=3,
                           n_jobs=-1)

# 训练模型（这实际上是在进行表格搜索）
grid_search.fit(X_train, y_train)

# 获取最佳参数
best_params = grid_search.best_params_
print(f"Best Parameters: {best_params}")

# 使用最佳参数重新训练模型
best_model = CatBoostRegressor(**best_params, loss_function='MAE', verbose=False)
best_model.fit(X_train, y_train, eval_set=(X_test, y_test), early_stopping_rounds=50)

# 预测测试集
y_pred = best_model.predict(X_test)

# 保存模型
best_model.save_model("catboost_model.cbm")

# 计算MAE
mae = mean_absolute_error(y_test, y_pred)
print(f'Mean Absolute Error: {mae}')